import torch
from torch import nn
import torch.nn.functional as F


class SiluAndMul(nn.Module):
    """
    模块功能
    SiluAndMul 模块实现的操作是：接收一个张量 x，先按最后一维（-1 维度 ）切分成两个形状相同的张量 x 和 y，然后对切分后的第一个张量 x 应用 SiLU（Sigmoid Linear Unit）激活函数，再将激活后的结果与第二个张量 y 逐元素相乘，最终返回相乘后的结果 ，公式可简单表示为 output = SiLU(x_part) * y_part ，常可用于自定义网络结构中特殊的激活与运算组合场景，比如在一些模型对中间特征进行特殊变换时可能会用到。
    实现细节
    初始化：__init__ 方法继承父类 nn.Module 的初始化逻辑，自身没有额外需要初始化的参数，因为该模块主要是定义前向传播逻辑。
    前向传播 forward 方法：
    x.chunk(2, -1)：将输入张量 x 在最后一个维度上分割成 2 个张量，假设 x 最后一维的大小是 2n，那么分割后得到的 x 和 y 最后一维大小会是 n ，这样切分是为了后续分别进行激活和相乘操作。
    F.silu(x)：对切分后的第一个张量 x 应用 SiLU 激活函数，SiLU 函数公式为 SiLU(x) = x * sigmoid(x) ，它是一种平滑的激活函数，在一些模型中相比 ReLU 等能带来更好的效果 。
    F.silu(x) * y：将激活后的结果与切分得到的第二个张量 y 进行逐元素相乘，得到最终输出。
    @torch.compile 装饰器：用于启用 PyTorch 的编译优化功能，它会尝试对 forward 方法的计算图进行优化，比如融合一些运算、优化内存使用、进行常量折叠等，从而可能提升该模块在训练或推理时的执行速度 。不过需要注意，torch.compile 也有一些限制和注意事项，比如某些复杂的操作可能无法很好编译优化，或者编译本身需要一定额外时间等，但在合适的场景下能有效加速计算。
    可能的应用场景
    这种模块可以用于构建自定义的神经网络层，例如在设计一些创新的 Transformer 变体、卷积神经网络中的特殊 bottleneck 结构时，当需要先对特征进行切分，一部分做激活另一部分参与相乘来调整特征分布，就可以嵌入 SiluAndMul 这样的模块。比如在一些追求高效计算和特殊特征变换的轻量级模型研发中可能会派上用场 。

    总体而言，SiluAndMul 是一个简洁但有特定功能的自定义模块，结合了切分、激活和逐元素相乘操作，借助 torch.compile 还有潜力提升执行效率，可灵活用于自定义网络架构搭建。
    """
    def __init__(self):
        super().__init__()

    @torch.compile
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, y = x.chunk(2, -1)
        return F.silu(x) * y
